# -*- coding: utf-8 -*-

# File Name : base_inference
# Description :
# Author : marje
# date : 2022/6/5
# Change Activity : 2022/6/5
from abc import abstractmethod

import torch

from . import is_key_right_
from ... import PipLine, join_relative_path, logger
from ...cv.misc.img_io import read_img
from ...dl.modules import load_dl_module_registry


class BaseInference:
    def __init__(self,
                 data_preprocess_pipline_cfg=None,
                 model_class=None,
                 model_config=None,
                 model_relative_path='',
                 save_step_detail=False,
                 model_input_key='img',
                 inference_result_key='inference_result'
                 ):
        if model_config is None:
            model_config = {}
        self.save_step_detail = save_step_detail
        self.model_class = model_class
        self.model_input_key = model_input_key
        self.model_config = model_config
        self.inference_result_key = inference_result_key
        self.model_relative_path = model_relative_path
        self.data_preprocess_pipline_cfg = BaseInference._format_data_preprocess_pipline_cfg(data_preprocess_pipline_cfg)
        self.data_preprocess_pipline = self._build_preprocess_pipline()
        self._model = self._init_model()
        self.post_process_pipline = self._build_post_process_pipline()

    @abstractmethod
    def _build_post_process_pipline(self):
        """
        build post process pipline.

        Returns:
            PipLine
        """
        pass

    @staticmethod
    def _get_registry():
        """
        Load default registry.
        Returns:

        """
        registry_lst = []
        registry = load_dl_module_registry()
        registry_lst.append(registry)
        return registry_lst

    def _build_preprocess_pipline(self):
        """
        Build pipline

        Returns:
            PipLine
        """
        registry_lst = BaseInference._get_registry()
        pipline = PipLine(self.data_preprocess_pipline_cfg, registry_lst, save_step_detail=self.save_step_detail)
        return pipline
        pass

    def _init_model(self):
        """
        Init model.
        Returns:
            model
        """
        model = self.model_class(**self.model_config)
        try:
            if is_key_right_:
                state_dict = torch.load(join_relative_path(self.model_relative_path))
                model.load_state_dict(state_dict)
        except Exception as e:
            logger.error('Failed to load model:' + str(e))
        return model.cuda().eval()

    @staticmethod
    def _format_data_preprocess_pipline_cfg(data_preprocess_pipline_cfg):
        """
        Format config.
        Args:
            data_preprocess_pipline_cfg:

        Returns:

        """
        if data_preprocess_pipline_cfg is None:
            data_preprocess_pipline_cfg = [dict(type='ImageToTensor', input_key='img', output_key='img'),
                                           dict(type='TensorToDevice', input_key='img', output_key='img', device_name='cuda:0'),
                                           dict(type='ToDataType', input_key='img', output_key='img', dtype=torch.float32)]
        return data_preprocess_pipline_cfg

    def inference_pipline(self, data):
        """

        Args:
            data:  eg:dict(img=img, device='cuda:0', to_dtype=torch.float32)

        Returns:

        """
        data = self.data_preprocess_pipline(data)
        with torch.no_grad():
            data = self._get_model_prediction(data)
        if self.post_process_pipline is not None:
            data = self.post_process_pipline(data)
        return data

    def _get_model_prediction(self, data):
        """
        get model prediction.

        Args:
            data:

        Returns:

        """
        img = data[self.model_input_key]
        res = self._model(img.clone())
        data[self.inference_result_key] = res
        return data

    def inference(self, img=None, path=None):
        """
        inference

        Args:
            img:
            path:

        Returns:

        """
        if img is None:
            img = read_img(path)
        data = {'img': img, 'path': path, 'original_img': img}
        data = self.inference_pipline(data)
        return data
